from audioop import reverse
import numpy as np
from collections import OrderedDict
from keras.datasets import mnist
import matplotlib.pylab as plt

'''
OrderedDict:有序字典，有序指的是它可以记住向字典里添加元素的顺序
'''


def softmax(x):
    if x.ndim == 2:
        x = x.T
        x = x - np.max(x, axis=0)
        y = np.exp(x) / np.sum(np.exp(x), axis=0)
        return y.T

    x = x - np.max(x)  # 溢出对策
    return np.exp(x) / np.sum(np.exp(x))


def cross_entropy_error(y, t):
    if y.ndim == 1:
        t = t.reshape(1, t.size)
        y = y.reshape(1, y.size)

    # 监督数据是one-hot-vector的情况下，转换为正确解标签的索引
    if t.size == y.size:
        t = t.argmax(axis=1)

    batch_size = y.shape[0]
    return -np.sum(np.log(y[np.arange(batch_size), t] + 1e-7)) / batch_size


# relu层的实现
class Relu:
    def __init__(self):
        self.mask = None

    def forward(self, x):
        # 前向传播小于或等于0的变为0
        # 大于0的原封不动
        self.mask = (x <= 0)
        out = x.copy()
        out[self.mask] = 0
        return out

    def backward(self, dout):
        # 反向传播小于或等于0的变为0
        # 大于0的乘以导数1也就是原封不动
        dout[self.mask] = 0
        dx = dout
        return dx


# Sigmoid层的实现
class Sigmoid:
    def __init__(self):
        self.out = None

    def forward(self, x):
        out = 1 / (1 + np.exp(-x))
        self.out = out
        return out

    def backward(self, dout):
        dx = dout * self.out * (1.0 - self.out)
        return dx


# Affine层的实现
class Affine:
    def __init__(self, W, b):
        self.W = W
        self.b = b
        self.x = None
        self.dW = None
        self.db = None

    def forward(self, x):
        self.x = x
        out = np.dot(x, self.W) + self.b
        return out

    def backward(self, dout):
        dx = np.dot(dout, self.W.T)
        self.dW = np.dot(self.x.T, dout)
        self.db = np.sum(dout, axis=0)
        return dx


# Softmax-with-Loss层的实现
class SoftmaxWithLoss:
    def __init__(self):
        self.loss = None
        self.y = None
        self.t = None

    def forward(self, x, t):
        self.t = t
        self.y = softmax(x)
        self.loss = cross_entropy_error(self.y, self.t)
        return self.loss

    def backward(self, dout=1):
        batch_size = self.t.shape[0]
        dx = (self.y - self.t) / batch_size
        return dx


# 两层网络结构
class TwoLayerNet:
    def __init__(self, input_size, hidden_size, output_size, weight_init_std=0.01):
        # 初始化权重
        self.params = {}
        self.params['W1'] = weight_init_std * np.random.randn(input_size, hidden_size)
        self.params['b1'] = np.zeros(hidden_size)
        self.params['W2'] = weight_init_std * np.random.randn(hidden_size, output_size)
        self.params['b2'] = np.zeros(output_size)

        # 生成层
        self.layers = OrderedDict()
        self.layers['Affine1'] = Affine(self.params['W1'], self.params['b1'])
        self.layers['Relu1'] = Relu()
        self.layers['Affine2'] = Affine(self.params['W2'], self.params['b2'])

        self.LastLayer = SoftmaxWithLoss()

    def predict(self, x):
        for layer in self.layers.values():
            x = layer.forward(x)
        return x

    # 损失函数
    def loss(self, x, t):
        y = self.predict(x)
        return self.LastLayer.forward(y, t)

    # 计算梯度
    def gradient(self, x, t):
        # 前向传播
        self.loss(x, t)
        # 反向传播
        # 反向传播必须从最后一层往回调用backward函数
        dout = 1
        dout = self.LastLayer.backward(dout)
        layers = list(self.layers.values())
        layers.reverse()
        for layer in layers:
            dout = layer.backward(dout)

        # 反向传播完后各个层的权重和偏置的梯度也就生成
        grads = {}
        grads['W1'] = self.layers['Affine1'].dW
        grads['b1'] = self.layers['Affine1'].db
        grads['W2'] = self.layers['Affine2'].dW
        grads['b2'] = self.layers['Affine2'].db

        return grads

    # 计算准确率
    def accuracy(self, x, t):
        y = self.predict(x)
        y = np.argmax(y, axis=1)
        if t.ndim != 1:
            t = np.argmax(t, axis=1)
        accuracy = np.sum(y == t) / float(x.shape[0])
        return accuracy


if __name__ == '__main__':
    (x_train, t_train), (x_test, t_test) = mnist.load_data()
    net = TwoLayerNet(input_size=784, hidden_size=50, output_size=10)

    train_loss_list = []

    # 超参数
    iter_nums = 10000
    train_size = x_train.shape[0]
    batch_size = 100
    learning_rate = 0.1

    # 记录准确率
    train_acc_list = []
    test_acc_list = []
    # 平均每个epoch的重复次数
    iter_per_epoch = max(train_size / batch_size, 1)

    for i in range(iter_nums):
        # 小批量数据
        batch_mask = np.random.choice(train_size, batch_size)
        x_batch = x_train[batch_mask]
        t_batch = t_train[batch_mask]

        # 计算梯度
        # 误差反向传播法 计算很快
        grad = net.gradient(x_batch, t_batch)

        # 更新参数 权重W和偏重b
        for key in ['W1', 'b1', 'W2', 'b2']:
            net.params[key] -= learning_rate * grad[key]

        # 记录学习过程
        loss = net.loss(x_batch, t_batch)
        print('训练次数:' + str(i) + '    loss:' + str(loss))
        train_loss_list.append(loss)

        # 计算每个epoch的识别精度
        if i % iter_per_epoch == 0:
            # 测试在所有训练数据和测试数据上的准确率
            train_acc = net.accuracy(x_train, t_train)
            test_acc = net.accuracy(x_test, t_test)
            train_acc_list.append(train_acc)
            test_acc_list.append(test_acc)
            print('train acc:' + str(train_acc) + '   test acc:' + str(test_acc))

    print(train_acc_list)
    print(test_acc_list)

    # 绘制图形
    markers = {'train': 'o', 'test': 's'}
    x = np.arange(len(train_acc_list))
    plt.plot(x, train_acc_list, label='train acc')
    plt.plot(x, test_acc_list, label='test acc', linestyle='--')
    plt.xlabel("epochs")
    plt.ylabel("accuracy")
    plt.ylim(0, 1.0)
    plt.legend(loc='lower right')
    plt.show()



